from diffusers import StableDiffusionPipeline
import torch
from PIL import Image
import numpy as np
import io

class MangaColorizer:
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.pipe = StableDiffusionPipeline.from_pretrained(
            "CompVis/stable-diffusion-v1-4",
            torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
        ).to(self.device)

    def colorize(self, image: Image.Image, prompt: str = "anime style, vibrant colors"):
        # Convert sketch to RGB if grayscale
        if image.mode != 'RGB':
            image = image.convert('RGB')
            
        # Generate colored image
        result = self.pipe(
            prompt=prompt,
            image=image,
            strength=0.7,
            guidance_scale=7.5,
            num_inference_steps=50
        ).images[0]
        
        # Convert to bytes
        img_byte_arr = io.BytesIO()
        result.save(img_byte_arr, format='PNG')
        return img_byte_arr.getvalue()

# Singleton instance
colorizer = MangaColorizer()
